import sys
from ui.acquiredialog import Ui_AcquireDialog
import wmi

from PyQt4.QtGui import *
class AcquireDialog (QDialog, Ui_AcquireDialog):
    OutputPath='out.csv'
    def __init__(self, parent=None):
        QDialog.__init__(self, parent)
        self.setupUi(self)
        c=wmi.WMI()
        self.drives=[]
        self.cbxDrives=[]
        for physical_disk in c.Win32_DiskDrive ():
            for partition in physical_disk.associators ("Win32_DiskDriveToDiskPartition"):
                for logical_disk in partition.associators ("Win32_LogicalDiskToPartition"):
                    cbx = QCheckBox(self.groupBox)
                    cbx.setText(logical_disk.Caption+' '+
                                physical_disk.Caption + partition.Caption)
                    self.verticalLayout.addWidget(cbx)
                    self.drives.append(logical_disk.Caption)
                    self.cbxDrives.append(cbx)

    def accept(self):
        if (self.txtNtuserHive.text()=='' or self.txtSystemHive.text()==''):
            QMessageBox.warning(self, 'IFE', 'Please set registry path' )
            return
        import ExecBat
        drives=''
        for i in range(self.cbxDrives.__len__()):
            if self.cbxDrives[i].isChecked():
                drives += self.drives[i][0]
        ExecBat.run(self.txtSystemHive, self.txtNtuserHive.text(),
                    self.spinDays.value(),drives,self.OutputPath)
        self.done(QDialog.Accepted)
        
    def browseFile(self):
        filename = QFileDialog.getOpenFileName(parent=self, caption="Select Hive file")
        if (filename==None):
            return
        else :
            if self.sender()==self.btnNtuser:
                self.txtNtuserHive.setText(filename)
            else: self.txtSystemHive.setText(filename)
            

def main():
    app = QApplication(sys.argv)
    d = AcquireDialog();
    d.exec_()
#mainwindow.exe()
   

if __name__=='__main__':main()